package com.wuwei.elearning.config.factory;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.wuwei.elearning.config.AiConfig;
import com.wuwei.elearning.question.model.ChatMessage;
import com.wuwei.elearning.question.model.ChatResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * Azure OpenAI客户端实现
 */
@Slf4j
public class AzureOpenAiClient implements AiClient {

    private final AiConfig aiConfig;
    private final RestTemplate restTemplate;
    private static final Pattern DATA_PATTERN = Pattern.compile("data: (.*)");
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    public AzureOpenAiClient(AiConfig aiConfig) {
        this.aiConfig = aiConfig;
        this.restTemplate = new RestTemplate();
    }
    
    @Override
    public String streamChatCompletion(List<ChatMessage> messages, SseEmitter emitter, StringBuilder responseBuilder) {
        try {
            // 将消息转换为OpenAI格式（与OpenAI相同）
            List<Map<String, String>> openaiMessages = convertToOpenAiFormat(messages);
            
            // 创建请求体
            Map<String, Object> requestBody = createRequestBody(openaiMessages, true);
            
            // 设置请求头
            HttpHeaders headers = createHeaders();
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
            
            // Azure OpenAI API URL格式
            String url = buildApiUrl();
            
            log.info("发送流式请求到Azure OpenAI API: {}", url);
            log.info("请求头: {}", headers);
            log.info("请求体: {}", requestBody);
            
            try {
                // 使用ResponseExtractor处理流式响应
                restTemplate.execute(url, HttpMethod.POST, 
                    req -> {
                        req.getHeaders().addAll(headers);
                        req.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                        try {
                            String jsonBody = objectMapper.writeValueAsString(requestBody);
                            req.getBody().write(jsonBody.getBytes());
                        } catch (IOException e) {
                            log.error("写入请求体时出错", e);
                        }
                    },
                    new ResponseExtractor<Void>() {
                        @Override
                        public Void extractData(ClientHttpResponse response) throws IOException {
                            try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getBody()))) {
                                String line;
                                while ((line = reader.readLine()) != null) {
                                    if (line.isEmpty()) continue;
                                    
                                    if (line.startsWith("data:")) {
                                        Matcher matcher = DATA_PATTERN.matcher(line);
                                        if (matcher.find()) {
                                            String data = matcher.group(1).trim();
                                            if ("[DONE]".equals(data)) {
                                                break;
                                            }
                                            
                                            try {
                                                // 解析JSON响应
                                                Map<String, Object> responseMap = objectMapper.readValue(data, Map.class);
                                                Map<String, Object> choices = (Map<String, Object>) ((List<?>) responseMap.get("choices")).get(0);
                                                Map<String, Object> delta = (Map<String, Object>) choices.get("delta");
                                                
                                                if (delta.containsKey("content")) {
                                                    String content = (String) delta.get("content");
                                                    responseBuilder.append(content);
                                                    
                                                    ChatResponse partialResponse = new ChatResponse();
                                                    partialResponse.setContent(responseBuilder.toString());
                                                    partialResponse.setIsComplete(false);
                                                    emitter.send(partialResponse);
                                                }
                                            } catch (Exception e) {
                                                log.error("解析流式响应时出错: {}", e.getMessage(), e);
                                            }
                                        }
                                    }
                                }
                            } catch (Exception e) {
                                log.error("处理流式响应时出错", e);
                                ChatResponse errorResponse = new ChatResponse();
                                errorResponse.setContent("处理API响应时出错: " + e.getMessage());
                                errorResponse.setIsComplete(true);
                                try {
                                    emitter.send(errorResponse);
                                } catch (IOException ex) {
                                    log.error("发送错误响应时出错", ex);
                                }
                            }
                            return null;
                        }
                    });
                
                log.info("Azure OpenAI API流式请求处理完成");
                
            } catch (Exception e) {
                log.error("Azure OpenAI API请求失败: {}", e.getMessage(), e);
                throw e;
            }
            
            return null; // 返回null表示内容已经通过emitter发送
            
        } catch (Exception e) {
            log.error("调用Azure OpenAI API时出错", e);
            return "错误：" + e.getMessage();
        }
    }
    
    public String chatCompletion(List<ChatMessage> messages) {
        try {
            // 将消息转换为OpenAI格式
            List<Map<String, String>> openaiMessages = convertToOpenAiFormat(messages);
            
            // 创建请求体
            Map<String, Object> requestBody = createRequestBody(openaiMessages, false);
            
            // 设置请求头
            HttpHeaders headers = createHeaders();
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
            
            // Azure OpenAI API URL格式
            String url = buildApiUrl();
            
            log.info("发送非流式请求到Azure OpenAI API: {}", url);
            
            // 发送请求
            Map<String, Object> response = restTemplate.postForObject(url, requestEntity, Map.class);
            
            if (response != null && response.containsKey("choices") && ((List)response.get("choices")).size() > 0) {
                Map<String, Object> choice = (Map<String, Object>) ((List)response.get("choices")).get(0);
                Map<String, String> message = (Map<String, String>) choice.get("message");
                return message.get("content");
            }
            
            return "无法从Azure OpenAI获取响应";
            
        } catch (Exception e) {
            log.error("调用Azure OpenAI API时出错", e);
            return "错误：" + e.getMessage();
        }
    }
    
    @Override
    public String getModelId() {
        return aiConfig.getAzureDeploymentId(); // Azure使用deploymentId作为模型ID
    }
    
    /**
     * 将内部消息格式转换为OpenAI API所需格式
     */
    private List<Map<String, String>> convertToOpenAiFormat(List<ChatMessage> messages) {
        return messages.stream()
                .map(msg -> {
                    Map<String, String> map = new HashMap<>();
                    map.put("role", msg.getRole());
                    map.put("content", msg.getContent());
                    return map;
                })
                .collect(Collectors.toList());
    }
    
    /**
     * 创建请求体
     */
    private Map<String, Object> createRequestBody(List<Map<String, String>> messages, boolean stream) {
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("messages", messages);
        requestBody.put("max_tokens", aiConfig.getMaxTokens());
        requestBody.put("temperature", aiConfig.getTemperature());
        requestBody.put("stream", stream);
        return requestBody;
    }
    
    /**
     * 创建HTTP请求头
     */
    private HttpHeaders createHeaders() {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("api-key", aiConfig.getApiKey());
        return headers;
    }
    
    /**
     * 构建Azure OpenAI API URL
     */
    private String buildApiUrl() {
        return String.format("%s/openai/deployments/%s/chat/completions?api-version=%s",
                aiConfig.getEndpoint(),
                aiConfig.getAzureDeploymentId(),
                aiConfig.getAzureApiVersion());
    }
} 